'''
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

Created on Jul 25, 2009

Version: 

@author: bryan
'''

import JT
import os
import numpy as np

class Volume(JT.vtk.vtkVolume):
    """
    Setup a volume  for rendering.
    """
    def __init__(self):

        #TODO:Allow the user to modify/customize the transfer functions to change how the DRR is rendered!!

        # The CTF setPoint will allow the brightness to be changed by adjusting this variable
        # Default values chosen from 2D3D gold standard data.
        self._ctfSetPoint = 86
        self._otfThreshold = 43
        self._otfLevel = 0.34
        self._rgb = [.0,.0,.0]

        #Default settings for the opacity transfer function,
        self._opacityTransferFunction = JT.vtk.vtkPiecewiseFunction()
        self._opacityTransferFunction.AddPoint(0.0, 0.0)
        self._opacityTransferFunction.AddPoint(self._otfThreshold, 0.0)
        self._opacityTransferFunction.AddPoint(255, self._otfLevel)

        # Create transfer mapping scalar value to color
        self._colorTransferFunction = JT.vtk.vtkColorTransferFunction()
        self._colorTransferFunction.AddRGBPoint(0.0, 0.0, 0.0, 0.0)
        self._colorTransferFunction.AddRGBPoint(self._ctfSetPoint, 1.0, 1.0, 1.0)
        self._colorTransferFunction.AddRGBPoint(255, 1.0, 1.0, 1.0)

        # The property describes how the data will look
        self._volumeProperty = JT.vtk.vtkVolumeProperty()

        # Setup default properties
        self._volumeProperty.SetColor(self._colorTransferFunction)
        self._volumeProperty.SetScalarOpacity(self._opacityTransferFunction)
        self._volumeProperty.ShadeOn()
        self._volumeProperty.SetInterpolationTypeToLinear()
        self.SetProperty(self._volumeProperty)

        #This is the key part for volumer rendering TextureMapper = fast DRR
        self._volumeMapper = JT.vtk.vtkVolumeTextureMapper3D()
        self.SetMapper(self._volumeMapper)

        self._maskReader = JT.vtk.vtkMetaImageReader()
        self._mask = JT.vtk.vtkImageMask()

        self._volImageData = None
        self._structuredPointsReader = JT.vtk.vtkStructuredPointsReader()
        self._metaImageReader = JT.vtk.vtkMetaImageReader()
        
        # Keep track of which body is actively being operated on
        self._activeBody = 0
        self.bodyCenters = []
        
        
    
    def UseRayCastMapper(self):
        """
        This will result in slower, but [potentially] more accurate rendering.
        """
        rcFunction = JT.vtk.vtkVolumeRayCastCompositeFunction()
        self._volumeMapper = JT.vtk.vtkVolumeRayCastMapper()
        self._volumeMapper.SetVolumeRayCastFunction(rcFunction)
        self.SetMapper(self._volumeMapper)
        if self._volImageData == None:
            print 'Image data not set'
            return
        self._volumeMapper.SetInput(self._volImageData)
        
    
    def UseTextureMapper(self):
        """
        This will result in faster, but [potentially] less accurate rendering.
        """
        self._volumeMapper = JT.vtk.vtkVolumeTextureMapper3D()
        self.SetMapper(self._volumeMapper)
        self._volumeMapper.SetInput(self._volImageData)

    def SetInput(self,vtkImageData):
        """
        Set the input image data for the mapper.
        """
        self._volumeMapper.SetInput(vtkImageData)

    def SetImageMask(self,maskPath):
        """
        Set the image mask.
        """
        if maskPath != None:
            self._maskReader.SetFileName(maskPath)
            self._maskReader.Update()
            self._mask.SetMaskInput(self._maskReader.GetOutput())
            self._mask.SetImageInput(self._volImageData)
            self._volumeMapper.SetInput(self._mask.GetOutput())
        else:
            self._volumeMapper.SetInput(self._volImageData)


    def SetVolumeFilename(self,volume_path):
        """
        Specify the path to a 3D volume and add to mapper.
        """
        # Select the correct reader based on the filename extension
        ext = volume_path[volume_path.rfind('.')+1:].lower()
        # logging.debug('Volume to be opened: %s' % volume_path)
        if (ext == 'mha'or ext == 'mhd'):
            self._reader = self._metaImageReader
        elif ext == 'vtk':
            self._reader = self._structuredPointsReader
        else:
            logging.debug("No reader has been implemented for %s type files" % ext)
            return

        self._reader.SetFileName(volume_path)
        self._volImageData = self._reader.GetOutput()
        
        # Set the image as the input to the shrinker, which is connected to the 
        # volume mapper:
        self.SetInput(self._reader.GetOutput())

        # if the user previous set custom display settings for this volume, use them
        self.OpenDisplaySettings(volume_path)
    
    
    def AddBodyCenter(self,center):
        """
        Set the rotational center of the vertebral body to register.
        
        Setting the center implies that a mask has been applied around the 
        region of interest.
        
        The centers of rotation for the vertebra are fixed.
        These are used to determined the initial position for
        each registration using the start positions 
        (which are the offsets from the ground truth), provided
        in a separate file (startpositions_200_ordered.txt)
        
        vert1 run1

        (x-127.5)*0.868=shiftx
        (y-127.5)*0.868=shifty
        (z-127.5)*0.868=shiftz
        
        centers in 3DRX voxels:
        b1 132 99 113
        b2 132 130 119
        b3 136 159 120
        
        centers in mm in world coordinates:
        b1 3.906 -24.738 -12.586
        b2 3.906 2.17 -7.378
        b3 7.378 27.342 -6.51
        ====================================================
        vert2 run4
        
        (x-127.5)*0.52108=shiftx
        (y-127.5)*0.52108=shifty
        (z-127.5)*0.52108=shiftz
        
        centers in 3DRX voxels:
        b1 143 47 99
        b2 135 87 94
        b3 127 129 87
        b4 117 171 78
        b5 108 214 74
        
        centers in mm in world coordinates:
        b1 8.07674 -41.94694 -14.85078
        b2 3.9081 -21.10374 -17.45618
        b3 -0.26054 0.78162 -21.10374
        b4 -5.47134 22.66698 -25.79346
        b5 -10.16106 45.07342 -27.87778
        """
        self.bodyCenters.append(center)
        self._activeBody = len(self.bodyCenters)-1
    
    def GetBodyCenter(self):
        return self.bodyCenters[self._activeBody]
    
    def GetVTKMatrixFromNumpyArray(self,npArray):
        """
        Return a VTK matrix from a numpy 4x4 array.
        """
        if (npArray.shape != (4,4)):
            print "Error: numpy array is not 4x4, cannot convert to VTK matrix"
            return
        
        vtkMatrix = JT.vtk.vtkMatrix4x4()
        
        for row in range(npArray.shape[0]):
            for col in range(npArray.shape[1]):
                vtkMatrix.SetElement(row,col,npArray[row][col])
        
        return vtkMatrix
    
    def GetNumpyArrayFromVTKMatrix(self,vtkMatrix):
        """
        Return a numpy array from a VTK Matrix
        """
        try:
            # Make sure that the user suppies a valid VTK matrix
            # This might also work if a transform is given
            vtkMatrix.GetClassName() == 'vtkMatrix4x4'
        except:
            print "Input is not recognized as a 'vtkMatrix4x4'"
            return
        
        npArray = np.eye(4)
        
        for row in range(npArray.shape[0]):
            for col in range(npArray.shape[1]):
                npArray[row][col] = vtkMatrix.GetElement(row,col,)
        
        return npArray
            
   
    def SetPose(self,pose):
        """
        Transform the volume according to the pose = Tx, Ty, Tz, Rx, Ry, Rz 
        Translations are in mm and rotations are in radians.
        
        this implementation is consistent with startpositions_200_ordered.txt
        
        This function will:
            1. shift the volume -center (so that rotation is about the center)
            2. rotate/translate the volume according to pose
            3. shift back by +center
            
        See Data_Help.txt for explanation of transformation process. 
        """
        # Top is TOriginalPositionWCS, the original position of the volume 
        # this should is specified during calibration (from sta file)
        # To set the start position, the volume will first be reset to the original
        # transform 
        self.SetTransform(self._originalTransform)
        transform = self.GetTransformFromPose(pose)
        startTransform = self.DoTransform(transform)
        self.SetTransform(startTransform)
        self._startTransform = startTransform
                    
    
    def SetOriginalTransform(self,transform):
        """
        Specify the original transform of the volume (4x4 numpy array).
        Only call this ONCE with starting postition of the volume.
        """
        # logging.debug('Volume original transform has been set to %s' % transform)
        self._originalTransform = transform
        # Set to User Matrix
        self.SetTransform(transform)
        
        
    def SetTransform(self,transform):
        """
        Specify the current transform of the volume (4x4 numpy array)
        this function will convert to a VTK matrix and update the current 
        user matrix of the volume. 
        The transform should be about the volumes center of rotation (not origin)
        """
        vtkMatrix = self.GetVTKMatrixFromNumpyArray(transform)
        self.SetUserMatrix(vtkMatrix)
        
    
    def DoTransform(self,transform):
        """
        Apply Transform after shifting the volume to the center of rotation. 
        This transform is concatenated to the current transform of the volume.
        This function will return a 4x4 numpy array.
        
        If SetStartPosition is called, it will reset the current pose to be 
        the original position of the volume before calling this function.
        """
        # Top is the original transform of the volume.
        Top = self._originalTransform
        # Tco is TTranscenterInv ( vertebral .bodyCenters center to volume origin)
        Tco = np.eye(4)
        Tco[0:3,3] = -np.asarray(self.bodyCenters[self._activeBody],'f')
        # Toc is TTranscenter (volume origin to vertebral body center)
        Toc = np.eye(4)
        Toc[0:3,3] = np.asarray(self.bodyCenters[self._activeBody],'f')
        # Tnp is new pose of the volume in WCS
        Tnp = np.dot(Toc,np.dot(transform,np.dot(Tco,Top)))
        return Tnp
       
  
    def GetResidualPose(self,transform1,transform2):
        """
        Return the pose that will transform the transform1->transform2.
        """
        # Top is TOriginalPositionWCS, the original position of the volume 
        # this should is specified during calibration (from sta file)
        Top = self._originalTransform
        # Tco is TTranscenterInv ( vertebral body center to volume origin)
        Tco = np.eye(4)
        Tco[0:3,3] = -np.asarray(self.bodyCenters[self._activeBody],'f')
        #T2shifted is transfrom2 shifted to center of body
        T2shifted = np.dot(Tco,transform2)
        #T1shifted is transform1 shifted to center of body)
        T1shifted = np.dot(Tco, transform1)
        #T1shiftedInverse is TOriginalErrorPositionInv (Original error postition invoerse)
        T1shiftedInverse = np.linalg.inv(T1shifted)
        # Tdiff is residual pose; the pose that will transform the original volume to the given transform
        Tdiff = np.dot(T2shifted,T1shiftedInverse)
        pose = self.GetPoseFromTransform(Tdiff)
        return pose
    
    def GetTransformFromPose(self,pose):
        """
        Given an arbitrary pose (6 element vector), return a transform (4x4 numpy array)
        derived from T=Rz*Ry*Rx.
        """
        Tx,Ty,Tz,Rx,Ry,Rz = pose
        
        # transform is TStartTransformation (body center to start pose)
        transform = np.eye(4)
        transform[0,0] = np.cos(Ry)*np.cos(Rz)
        transform[0,1] = -np.cos(Rx)*np.sin(Rz) + np.cos(Rz)*np.sin(Rx)*np.sin(Ry)
        transform[0,2] = np.sin(Rx)*np.sin(Rz) + np.cos(Rx)*np.cos(Rz)*np.sin(Ry)
        transform[0,3] = Tx
        transform[1,0] = np.cos(Ry)*np.sin(Rz)
        transform[1,1] = np.cos(Rx)*np.cos(Rz) + np.sin(Rx)*np.sin(Ry)*np.sin(Rz)
        transform[1,2] = -np.cos(Rz)*np.sin(Rx) + np.cos(Rx)*np.sin(Ry)*np.sin(Rz)
        transform[1,3] = Ty
        transform[2,0] = -np.sin(Ry)
        transform[2,1] = np.cos(Ry)*np.sin(Rx)
        transform[2,2] = np.cos(Rx)*np.cos(Ry)
        transform[2,3] = Tz
        
        return transform
        
    
    def GetPoseFromTransform(self,transform):
        """
        Given a 4x4 transform in numpy, return a 6 element vector 
        (Tx,Ty,Tz,Rx,Ry,Rz).  Rotation order: T=Rz*Ry*Rx [x->y->z]
        """
        Tx = transform[0][3]
        Ty = transform[1][3]
        Tz = transform[2][3]
        Ry = -np.arcsin(transform[2][0])
        Rx = np.arcsin(transform[2][1]*1/np.cos(Ry))
        Rz = np.arcsin(transform[1][0]*1/np.cos(Ry)) 
        
        pose=np.array([Tx,Ty,Tz,Rx,Ry,Rz])
        
        return pose
    
    def GetPose(self): 
        """
        Return the current 6 element vector (Tx,Ty,Tz,Rx,Ry,Rz)
        That describes the position and orientation of the volume.
        """
        # transform is the current transformation of the volume
        transform = self.GetNumpyArrayFromVTKMatrix(self.GetUserMatrix())
        pose = self.GetPoseFromTransform(transform)
        return pose     
    
    def GetPoseFromOriginalToCurrent(self): 
        """
        Return the current 6 element vector (Tx,Ty,Tz,Rx,Ry,Rz)
        That describes the postition and orientation of the volume
        with respect to the original pose of the volume.
        """
        # transform is the current transformation of the volume
        transform = self.GetNumpyArrayFromVTKMatrix(self.GetUserMatrix())
        pose = self.GetResidualPose(self._originalTransform,transform)
        return pose 


    def OpenDisplaySettings(self,volume_path):
        """
        Given the path to the volume that is being opened for rendering, find
        (if it exists) a display settings file and read the settings in to be
        used for volume rendering.
        """
        # Try to open volume dispay settings (a file with same name as volume,
        # with a .disp extension)
        (basePath,volName) = os.path.split(volume_path)
        shortVolName = volName[:volName.rfind('.')]
        dispSettingsName = os.path.join(basePath,shortVolName+'.disp')
        self._dispSettingsName = dispSettingsName

        try:
            dispFile = open(dispSettingsName)
            line1=dispFile.readline().split()
            line2=dispFile.readline().split()

            # Make sure the file contains the expected info
            if line1[0] == "opacityLevel:":
                opacityLevel = float(line1[1])
            else:
                print "Error: %s is not the expected variable name" % line1[0]

            if line2[0] == "opacityThreshold:":
                opacityThreshold = float(line2[1])
            else:
                print "Error: %s is not the expected variable name" % line1[0]

        except:
            print "Error opening display settings %s " % dispSettingsName
            # logging.debug("Could not find display settings file, one will be created if the default values are changed")

        self.SetOpacityLevel(opacityLevel)
        self.SetOpacityThreshold(opacityThreshold)
        #self.UpdateVolumeProperties()

    def SetVolumeColor(self,rgb=[0,0,0]):
        """
        Change the color of the rendered volume so that it can be blended
        together with the fixed image for better visualization

        Default is white on black [1,1,1]
        """
        self._rgb = rgb
        self._colorTransferFunction.AddRGBPoint(0.0, 0.0, 0.0, 0.0)
        self._colorTransferFunction.AddRGBPoint(self._ctfSetPoint, rgb[0], rgb[1], rgb[2])
        self._colorTransferFunction.AddRGBPoint(255, rgb[0], rgb[1], rgb[2])
        #self.UpdateVolumeProperties()

    def SetBrightness(self,ctfSetPoint):
        """
        A method to allow the brightness of the color transfer function to be set.
        The brightness is set by specifying the midpoint on the ramp function,
        the low point is fixed at 0 and the high point is fixed at 255.

        setPoint should be a scalar between 0 and 255.
        """
        # First Remove the existing point
        index = self._colorTransferFunction.RemovePoint(self._ctfSetPoint)
        if index == -1:
            print "Error setting color point."
            return

        # Save the new point for use later
        self._ctfSetPoint = ctfSetPoint
        self._colorTransferFunction.AddRGBPoint(self._ctfSetPoint, self._rgb[0], self._rgb[1], self._rgb[2])

    def SetOpacityThreshold(self,otfThreshold):
        """
        method to allow the transparency of the opacity transfer function to be set.
        The Opacity is set by specifying the threshold on the ramp function,
        the low point is fixed at 0 and the high point is fixed at 255.

        setThreshold should be an integer between 0 and 255.
        """
        # First Remove the existing point
        index = self._opacityTransferFunction.RemovePoint(self._otfThreshold)
        if index == -1:
            print "Error setting opacity threshold."
            return

        # Save the new point for use later
        self._otfThreshold = otfThreshold
        self._opacityTransferFunction.AddPoint(self._otfThreshold, 0.0)

    def SetOpacityLevel(self,otfLevel):
        """
        method to allow the transparency of the opacity transfer function to be set.
        This function changes the maximum opacity of the volume.

        otfValue should be a float between 0 and 1.
        """
        # Check to make sure the values is between 0 and 1
        if (otfLevel < 0) or (otfLevel > 1):
            print "Error setting opacity value."
            return

        # Save the new point for use later
        self._otfLevel = otfLevel
        # The point is currently fixed at 255,should this be adjustable
        self._opacityTransferFunction.AddPoint(255, self._otfLevel)
        
        

class SurfaceModel(JT.vtk.vtkActor):
    """
    Create an actor that can be used an isosurface.
    """
    def __init__(self):
        self._mapper=JT.vtk.vtkPolyDataMapper()
        self._maskReader = JT.vtk.vtkMetaImageReader()
        self._mask = JT.vtk.vtkImageMask()
        self._surfaceModel = None
        
    def SetInput(self,data):
        """
        Set the input image data for the mapper.
        """
        self._mapper.SetInput(data)
        
    def SetVolumeFilename(self,volume_path):
        """
        Specify the path to a 3D volume and add to mapper.
        """
        # Select the correct reader based on the filename extension
        ext = volume_path[volume_path.rfind('.')+1:].lower()
        # logging.debug('Volume to be opened: %s' % volume_path)
        if (ext == 'mha'or ext == 'mhd'):
            reader = JT.vtk.vtkMetaImageReader()
        elif ext == 'vtk':
            reader = JT.vtk.vtkStructuredPointsReader()
        else:
            logging.debug("No reader has been implemented for %s type files" % ext)
            return

        reader.SetFileName(volume_path)
        #self.SetInput(reader.GetOutput())
        self._volImageData = reader.GetOutput()


    def CreateIsoSurface(self):
        """
        Create an iso surface.
        
        TODO: allow the user to set the threshold for marching cubes algorithm.
        """
        
        if self._maskPath != None:
            # Volume will only display masked region
            self._displayData = self._mask.GetOutput()
        else:
            # Volume will display whole extent
            self._displayData = self._volImageData.copy()
            
        marching=JT.vtk.vtkImageMarchingCubes()
        marching.SetInput(self._displayData)
        marching.SetValue(0,80)
        self._mapper.SetInput(marching.GetOutput())
        self.SetMapper(self._mapper)
        prop=self.GetProperty()
        prop.SetColor(.2,.2,.2)
        self._surfaceModel = marching.GetOutput()
                
    def SetImageMask(self,maskPath):
        """
        Set the image mask.  To turn mask off, set maskPath to None.
        """
        if maskPath != 0:
            self._maskPath = maskPath
            self._maskReader.SetFileName(self._maskPath)
            self._maskReader.Update()
            self._mask.SetMaskInput(self._maskReader.GetOutput())
            self._mask.SetImageInput(self._volImageData)
        else:
            print "No mask provided"
            
    def AddBodyCenter(self,center):
        """
        See documentation in JT.Volume()
        """
        self.bodyCenters.append(center)
        
        
    def SetOriginalTransform(self,transform):
        """
        Specify the original transform of the volume (4x4 numpy array).
        Only call this ONCE with starting postition of the volume.
        """
        self._originalTransform = transform
        # Set to User Matrix
        self.SetTransform(transform)
    
    def SetTransform(self,transform):
        """
        Specify the current transform of the volume (4x4 numpy array)
        this function will convert to a VTK matrix and update the current 
        user matrix of the volume. 
        The transform should be about the volumes center of rotation (not origin)
        Using the TransformAboutCenter function can facilitate this.
        """
        vtkMatrix = self.GetVTKMatrixFromNumpyArray(transform)
        self.SetUserMatrix(vtkMatrix)
        
    def GetVTKMatrixFromNumpyArray(self,npArray):
        """
        Return a VTK matrix from a numpy 4x4 array.
        """
        if (npArray.shape != (4,4)):
            print "Error: numpy array is not 4x4, cannot convert to VTK matrix"
            return
        
        vtkMatrix = JT.vtk.vtkMatrix4x4()
        
        for row in range(npArray.shape[0]):
            for col in range(npArray.shape[1]):
                vtkMatrix.SetElement(row,col,npArray[row][col])
        
        return vtkMatrix
    
    def SaveAsSTL(self,filename):
        """
        """
        if self._surfaceModel != None:
            writer = JT.vtk.vtkSTLWriter()
            writer.SetInput(self._surfaceModel)
            writer.SetFileName(filename)
            writer.Write()
